import pickle
import numpy as np
import pandas as pd
import tensorly as tl

backend="cupy"
tl.set_backend(backend)
device = 'cuda'

K = 20

tlda = pickle.load(open('data/metoo_experiment/num_tops_10_alpha0_0.001_learning_rate_0.0005_theta_5.005_orthogonality_1000_initialize_first_docs_True_n_eigenvec_40_n_docs_0_no_online/tlda.obj', 'rb'))
factors_tlda = tlda.unwhitened_factors.get().T

vocab = np.array(pd.read_csv('data/metoo_experiment/vocab.csv')['words'])

top_words_tlda = {}
for i, f in enumerate(factors_tlda):
    top_tlda = np.argpartition(f,-K)[-K:]
    top_words_tlda[i] = list(vocab[top_tlda])

print(top_words_tlda)
df = pd.DataFrame(top_words_tlda)
df.to_csv('results/metoo_topics_tlda.csv')